import librosa
import math
import numpy as np
import sklearn
import os
import sys


# def getDataAndLabels(path):
#     """
#     得到全数据集
#     """
#     dirs = os.listdir(path)
#     #print(dirs)
#     # foldersPath = []
#     filesContainer = [[] for i in range(10)]

#     for dir in dirs:
#         tmpList = []
#         folderPath = os.path.join(path, dir)
#         files = os.listdir(folderPath)
#         for file in files:
#             filepath = os.path.join(folderPath, file)
#             #print(filepath)
#             # 加载数据
#             data, sampling_rate = librosa.load(path = filepath, sr = 44100, mono = True)
#             #print(data.shape)
#             assert data.shape == (88200,)
#             tmpList.append(data)
#         filesContainer[int(dir, 10)] = tmpList
#     return filesContainer

def getDataAndLabels(path):
    """
    得到全数据集
    """
    dirsf = os.listdir(path)
    print(dirsf)
    # foldersPath = []
    filesContainer = [[] for i in range(10)]

    for dirf in dirsf:
        # tmpList = []
        folderPathf = os.path.join(path, dirf)
        dirs = os.listdir(folderPathf)
        for dir in dirs:
            folderPath = os.path.join(folderPathf, dir)
            files = os.listdir(folderPath)
            for file in files:
                filepath = os.path.join(folderPath, file)
                print(filepath)
                # 加载数据
                data, sampling_rate = librosa.load(path=filepath, sr=44100, mono=True)
                # print(data.shape)
                # assert data.shape == (88200,)
                # tmpList.append(data)
                if data.shape == (88200,):
                    filesContainer[int(dir, 10)].append(data)
    minn = 10000
    for i in range(10):
        minn = min(minn, len(filesContainer[i]))
    for i in range(10):
        filesContainer[i] = filesContainer[i][:minn]
    return filesContainer


def devideAndWindow(input_signal, sr=44100, frame_length=882, shift_length=441, window_function="juxing"):
    """
    分帧，窗口
    """
    print(np.array(input_signal).shape)
    # shift_points = sr * shift_length
    # frame_points = sr * frame_length
    shift_points = shift_length
    frame_points = frame_length
    results = []
    for i in range(0, int(len(input_signal) - frame_points + 1), int(shift_points)):
        tmp = np.array(input_signal[i: i + int(frame_points)])
        if window_function == "haming":
            # print("haming")
            for j in range(len(tmp)):
                tmp[j] = tmp[j] * (0.54 - 0.46 * math.cos(2 * math.pi / (len(tmp) - 1)))
        elif window_function == "haining":
            # print("haining")
            for j in range(len(tmp)):
                tmp[j] = tmp[j] * (0.5 * (1 - math.cos(2 * math.pi / (len(tmp) - 1))))
        elif window_function != "juxing":
            raise Exception("Please select window function in (juxing, hanming, haining).")
        results.append(tmp)
    results = np.array(results)
    print(results.shape)
    return results


def wrapped(all_data, sr=44100, frame_length=882, shift_length=441, window_function="juxing"):
    """
    封装数据结构，也就是分帧后的全数据集结构，并将封装后数据返回
    """
    h, w = len(all_data), len(all_data[0])
    wrappedDataAfterWindow = [[] for i in range(10)]
    for i in range(h):
        w1 = len(all_data[i])
        for j in range(w1):
            afterWindow = devideAndWindow(all_data[i][j], sr, frame_length, shift_length, window_function)
            wrappedDataAfterWindow[i].append(afterWindow)

    return np.array(wrappedDataAfterWindow), np.array(all_data)


def activateFromOutside(path="./audio", sr=44100, frame_length=882, shift_length=441, window_function="juxing"):
    all_data = getDataAndLabels(path)
    print(np.array(all_data).shape)
    wrapped_data, all_data = wrapped(all_data, sr, frame_length, shift_length, window_function)
    # wrapped_data = None
    return wrapped_data, all_data


'''
if __name__ == "__main__":
    wrapped_data = activateFromOutside(window_function="haming")
    print(wrapped_data.shape)
    '''
